import sys
import getopt
import time

import pandas as pd
from icecream import ic

from difflib import SequenceMatcher as difsm


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


def job_stat(fun):
    def _call_(*argv, **kwgs):
        start_time = time.time()
        print('*** JOB START ***')
        fun(*argv, **kwgs)
        print('*** JOB FINISHED ***')
        print('. START    : %s' % (time.strftime("%H:%M:%S | %Y-%m-%d", time.localtime(start_time))))
        print('. FINISHED : %s' % (time.strftime("%H:%M:%S | %Y-%m-%d", time.localtime())))
        print('. TIME COST: %f s' % (time.time() - start_time))

    return _call_


def my_code_entry(ntab5225_275snp_zc208fp_hmp, query_snp_tsv):
    print('*** 程序运行需要一定的时间，请耐心等候 ***')
    # 简并碱基转换字典
    nt_dict = {'A': 'AA', 'C': 'CC', 'G': 'GG', 'T': 'TT',
               'R': 'AG', 'Y': 'CT', 'S': 'CG', 'W': 'AT',
               'K': 'GT', 'M': 'AC', 'N': 'NN'}

    df_ntab5225 = pd.read_csv(ntab5225_275snp_zc208fp_hmp, sep='\t', header=0, index_col='rs#')
    tgt_snp = list(df_ntab5225.index)  # 入选指纹图谱的全部275个SNP位点的rs#
    germ5225 = df_ntab5225.columns.tolist()[10:]  # 取出了全部种质资源的编号
    # ic(','.join(germ5225))

    df_query = pd.read_csv(query_snp_tsv, sep='\t', header=0, index_col='rs#')
    # ic(df_query)

    genetic_dist_mtx = []

    # 逐个种质资源遍历5225份种质资源的对应275个snp位点的指纹图谱记录
    for germ in germ5225:
        # ic(germ)
        dist_mtx = []

        # 逐行遍历用户输入的数据，只有在目标snp列表中的位点才会分析
        for idx, row in df_query.iterrows():
            # ic(df_query.loc)
            if idx in tgt_snp:
                # ic(idx)
                # 当前用户输入的snp是指纹图谱的snp
                # 取出用户数据的碱基，并从单字母转成双字母的ATCG类型
                usr_geno = nt_dict[df_query.loc[idx, 'Query']]
                hmp_geno = nt_dict[df_ntab5225.loc[idx, germ]]
                if usr_geno != 'NN' and hmp_geno != 'NN':
                    # 在均不为空缺的情况下
                    ratio = difsm(None, usr_geno, hmp_geno).quick_ratio()
                    # AA与AA，ratio 为 1.0
                    # AT与TA，ratio 为 1.0
                    # AA与TA，ratio 为 0.5
                    # AA与TT，ratio 为 0
                    dc = {'germplasm': germ,
                          'rs': idx,
                          'usr_geno': usr_geno,
                          'hmp_geno': hmp_geno,
                          'IBS': ratio * 2,
                          'GeneticDistance': 1.0 - ratio}
                    dist_mtx.append(dc)
        # 完成了全部计算，找出遗传距离最近的
        try:
            # 碰上缺失特别多的种质，有可能在所有位点上它都是缺失。len(dist_mtx)有等于0的可能。
            gdist = sum([dc['GeneticDistance'] for dc in dist_mtx]) / len(dist_mtx)
            genetic_dist_mtx.append((round(gdist, 4), germ,))
        except:
            # ic(germ)
            continue

    genetic_dist_mtx.sort(key=lambda x: x[0])
    # ic(genetic_dist_mtx[:10])
    # 输出分析结果
    print('排序\t品种\t遗传距离')
    for i, gdist in enumerate(genetic_dist_mtx[:100]):
        dist_value, germ = gdist
        print(f'{i + 1}\t{germ}\t{dist_value}')


# @job_stat
def main(argv=None):
    if argv is None:
        argv = sys.argv
    try:
        try:
            # 两个默认的数据文件
            # ntab5225_275snp_zc208fp_hmp = 'Ntab5225.275snp.zc208fp.hmp'
            ntab5225_275snp_zc208fp_hmp = r'..\anal_202209\zc208.pop5161.fingerprint.1.hmp'
            query_snp_tsv = 'query_snp.tsv'

            # 开始解析参数
            opts, args = getopt.getopt(argv[1:], 'hd:q:', ['help', 'database=', 'query='])
            # https://www.jianshu.com/p/a877e5b46b2d
            for opt_name, opt_value in opts:
                if opt_name in ('-d', '--database'):
                    ntab5225_275snp_zc208fp_hmp = opt_value

                if opt_name in ('-q', '--query'):
                    query_snp_tsv = opt_value

        except getopt.error as msg:
            raise Usage(msg)
        # more code, unchanged
        my_code_entry(ntab5225_275snp_zc208fp_hmp, query_snp_tsv)

    except Usage as err:
        print(err.msg, file=sys.stderr)
        print("for help use --help", file=sys.stderr)
        return 2


if __name__ == "__main__":
    sys.exit(main())  # https://www.artima.com/weblogs/viewpost.jsp?thread=4829
